package com.example.log.alarm.agent;

import net.bytebuddy.ByteBuddy;
import net.bytebuddy.agent.ByteBuddyAgent;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.NamedElement;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.method.ParameterList;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.ClassFileLocator;
import net.bytebuddy.dynamic.loading.ClassInjector;
import net.bytebuddy.matcher.ElementMatcher;
import net.bytebuddy.matcher.ElementMatchers;
import net.bytebuddy.pool.TypePool;
import net.bytebuddy.utility.JavaModule;
import org.springframework.util.StringUtils;

import java.lang.instrument.Instrumentation;
import java.util.*;
import java.util.stream.Collectors;

import static net.bytebuddy.matcher.ElementMatchers.nameContains;
import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;

/**
 * @program: spring-agent
 * @description:
 * @author: zt22328 赵儒定
 * @create: 2023-03-31 09:15
 **/
public class AgentTemplate {

    private List<AgentVisitBase> visits;

    private String[] spyClassNames;

    public AgentTemplate(List<AgentVisitBase> visits, String ... spyClassNames) {
        this.visits = visits;
        this.spyClassNames = spyClassNames;
    }

    public void start() throws ClassNotFoundException {
        Instrumentation install = ByteBuddyAgent.install();
        AgentBuilder agentBuilder = new AgentBuilder.Default();
        injection(install, agentBuilder, spyClassNames);

        agentBuilder
                .with(AgentBuilder.InitializationStrategy.NoOp.INSTANCE)
                .with(AgentBuilder.RedefinitionStrategy.REDEFINITION)
                .with(AgentBuilder.TypeStrategy.Default.REDEFINE)
                .ignore(ignore())
                .type(type())
                .transform(transformer())
                .with(new DefaultAgentListener())
                .installOnByteBuddyAgent();
    }

    public static String internalDelegate(String methodsInterceptor) {
        return methodsInterceptor + "_internal";
    }

    public void injection(Instrumentation install, AgentBuilder agentBuilder, String ... classNames) throws ClassNotFoundException {
        Map<String, byte[]> classesTypeMap = new HashMap<>();
        for (String className : classNames) {
            TypeDescription typeDescription = TypePool.Default.of(this.getClass().getClassLoader())
                    .describe(className).resolve();
            String internalDelegate = internalDelegate(className);
            classesTypeMap.put(internalDelegate,
                    new ByteBuddy()
                            .redefine(typeDescription, ClassFileLocator.ForClassLoader.of(this.getClass().getClassLoader()))
                            .name(internalDelegate)
                            .make()
                            .getBytes());
        }
        ClassInjector.UsingUnsafe.Factory factory = ClassInjector.UsingUnsafe.Factory.resolve(install);
        factory.make(null, null).injectRaw(classesTypeMap);
        agentBuilder.with(new AgentBuilder.InjectionStrategy.UsingUnsafe.OfFactory(factory));
        for (String key : classesTypeMap.keySet()) {
            agentBuilder.assureReadEdgeFromAndTo(install, Class.forName(key));
        }
    }

    public AgentBuilder.Transformer transformer() {
        return (builder, typeDescription, classLoader, javaModule)-> {
            String actualName = typeDescription.getActualName();
            Optional<AgentVisitBase> first = visits.stream().filter(e -> e.className.equals(actualName)).findFirst();
            if (!first.isPresent()) {
                return builder;
            }
            AgentVisitBase define = first.get();
            return builder.visit(Advice.to(define.getAdviceInterceptor()).on(m-> constructorOrMethodVisit(m, define)));
        };
    }

    private boolean constructorOrMethodVisit(MethodDescription m, AgentVisitBase define) {
        switch(define.interceptorType) {
            case METHOD:
                return methodVisit(m, define.getMethods());
            case CONSTRUCTOR:
                return constructorVisit(m, define.getMethods());
            case CONSTRUCTOR_AND_METHOD:
                return methodVisit(m, define.getMethods()) || constructorVisit(m, define.getMethods());
            default:
                return false;
        }
    }

    private boolean constructorVisit(MethodDescription m, List<AgentVisitBase.MethodDefinition> methods) {
        if (!m.isConstructor()) {
            return false;
        }
        List<AgentVisitBase.MethodDefinition> collect =
                methods.stream().filter(e -> StringUtils.isEmpty(e.getName())).collect(Collectors.toList());
        if (isEmpty(collect)) {
            return true;
        }
        for (AgentVisitBase.MethodDefinition method : collect) {
            if (method.isAllMethods() || (StringUtils.isEmpty(method) && paramTypesMatch(m.getParameters(), method))) {
                return true;
            }
        }
        return false;
    }

    private boolean methodVisit(MethodDescription m, List<AgentVisitBase.MethodDefinition> methods) {
        if (!m.isMethod()) {
            return false;
        }
        for (AgentVisitBase.MethodDefinition method : methods) {
            if (m.getActualName().equals(method.getName())
                    && (method.isAllMethods() || paramTypesMatch(m.getParameters(), method))) {
                return true;
            }
        }
        return false;
    }

    private boolean isEmpty(String str) {
        return str == null || str.length() <= 0;
    }

    private boolean paramTypesMatch(ParameterList<?> parameters, AgentVisitBase.MethodDefinition method) {
        if (method.isAllMethods()) {
            return true;
        }
        List<String> paramTypes = method.getParamTypes();
        if (isEmpty(parameters) && isEmpty(paramTypes)) {
            return true;
        }
        if (parameters.size() != paramTypes.size()) {
            return false;
        }
        int num = 0;
        for (int i = 0; i < parameters.size(); i++) {
            if (parameters.get(i).getType().getTypeName().equals(paramTypes.get(i))) {
                num ++;
            }
        }
        return parameters.size() == num;
    }

    public static boolean isEmpty(List list) {
        return list == null || list.isEmpty();
    }

    private ElementMatcher<? super TypeDescription> type() {
        List<String> names = classNameMatches();
        if (isEmpty(names)) {
            return ElementMatchers.none();
        }
        return ElementMatchers.namedOneOf(names.toArray(new String[names.size()]));
    }

    /**
     * 排除包
     * @return
     */
    private ElementMatcher<? super TypeDescription> ignore() {
        ElementMatcher.Junction<NamedElement> junction = nameStartsWith("net.bytebuddy.")
                .or(nameStartsWith("org.slf4j."))
                .or(nameStartsWith("org.groovy."))
                .or(nameContains("javassist"))
                .or(nameContains(".asm."))
                .or(nameContains(".reflectasm."))
                .or(nameStartsWith("sun.reflect"))
//                .or(nameStartsWith("java."))
                ;
        List<String> startsWiths = ignoreNameStartsWiths();
        if (startsWiths != null && !startsWiths.isEmpty()) {
            startsWiths.forEach(s->{
                junction.or(nameStartsWith(s));
            });
        }
        List<String> contains = ignoreNameContains();
        if (contains != null && !contains.isEmpty()) {
            contains.forEach(s->{
                junction.or(nameContains(s));
            });
        }
        return junction;
    }

    protected List<String> ignoreNameStartsWiths() {
        return Collections.EMPTY_LIST;
    }

    protected List<String> ignoreNameContains() {
        return Collections.EMPTY_LIST;
    }

    protected void withs(AgentBuilder agentBuilder) {

    }

    private List<String> classNameMatches() {
        if (isEmpty(this.visits)) {
            return Collections.EMPTY_LIST;
        }
        return this.visits.stream().map(AgentVisitBase::getClassName).collect(Collectors.toList());
    }



    public static class DefaultAgentListener extends AgentBuilder.Listener.Adapter {

        @Override
        public void onDiscovery(String typeName, ClassLoader classLoader, JavaModule module, boolean loaded) {
        }

        public void onError(String typeName, ClassLoader classLoader, JavaModule module, boolean loaded, Throwable throwable) {

        }

    }

}
